from typing import Optional, List, Any, Mapping
import requests
import ollama
from langchain_core.language_models.llms import LLM

class OllamaLLM(LLM):
    model: str = "deepseek-r1:8b"  # 默认模型名称
    tools: Optional[List[Any]] = None  # 声明 tools 字段

    def _llm_type(self) -> str:
        return "ollama"

    def bind_tools(self, tools: List[Any]) -> "OllamaLLM":
        self.tools = tools
        return self

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        messages = [
            {"role": "system", "content": "你是一个金融分析专家，回答要求精准且简明。"},
            {"role": "user", "content": prompt}
        ]
        # 使用 self.model 作为模型名称
        response = ollama.chat(model=self.model, messages=messages)
        return response['message']['content']

# 示例调用
if __name__ == '__main__':
    ollama_llm = OllamaLLM()
    print(ollama_llm._call("请解释一下量化交易的基本原理。"))
